import tensorflow as tf
from tensorflow import keras
print(tf.__version__) # 打印出TensorFlow的版本

# 设置一次送入GPU的量
batchsize = 32

# LeNet5网络模型
model = tf.keras.Sequential([
    # 第1层-卷积层 滤波器尺寸5*5,6个
    keras.layers.Conv2D(6,5),
    # 第2层-池化层,滤波器尺寸为2×2，步长为2
    keras.layers.MaxPooling2D(pool_size=2,strides=2),
    keras.layers.ReLU(),
    # 第3层-卷积层 滤波器尺寸5*5,16个
    keras.layers.Conv2D(16,5),
    # 第4层-池化层,滤波器尺寸为2×2，步长为2
    keras.layers.MaxPooling2D(pool_size=2,strides=2),
    keras.layers.ReLU(),
    keras.layers.Flatten(), # 拉平之后方便做全连接
    # 第5层-全连接层
    keras.layers.Dense(120,activation='relu'),
    # 第6层-全连接层
    keras.layers.Dense(84,activation='relu'),
    # 第7层-全连接层
    keras.layers.Dense(10,activation='softmax') # 最后输出10类,0-9的数字
])


model.build(input_shape=(batchsize,28,28,1)) # 28*28*1 1个通道

# 打印出网络结构
model.summary()

# 编译
model.compile(optimizer=keras.optimizers.Adam(),
              loss = keras.losses.CategoricalCrossentropy(),
              metrics=['accuracy'])

# 数据预处理
def preprocess(x,y):
    x = tf.cast(x,dtype=tf.float32) / 255
    x = tf.reshape(x,[-1,28,28,1]) #
    y = tf.one_hot(y,depth=10)
    return x,y

# 加载数据集
(x_train,y_train),(x_test,y_test) = tf.keras.datasets.mnist.load_data()

print('x_train=',type(x_train))
print('y_train=',type(y_train))

train_db = tf.data.Dataset.from_tensor_slices((x_train,y_train))
print('train_db=',type(train_db))
train_db = train_db.shuffle(10000)
train_db = train_db.batch(128)
train_db = train_db.map(preprocess)

test_db = tf.data.Dataset.from_tensor_slices((x_test,y_test))
print('test_db=',type(test_db))
test_db = test_db.shuffle(10000)
test_db = test_db.batch(128)
test_db = test_db.map(preprocess)

# 开始训练
model.fit(train_db,epochs=10)

# 在测试集上验证
model.evaluate(test_db)

# 保存训练的模型
model.save('./lenet5.h5')